import os
import logging
from split_function import split_rust_functions_in_file
from extract_component import extract_function_names, extract_static_declaration

def transform_unsafe_fn(rust_code: str) -> str:
    # 查找函数体的起始位置
    body_start = rust_code.find("{")
    if body_start == -1:
        raise ValueError("Function body not found")

    # 查找函数体的结束位置
    body_end = rust_code.rfind("}")
    if body_end == -1:
        raise ValueError("Function body not found")

    # 提取函数签名和函数体
    signature = rust_code[:body_start].strip()
    body = rust_code[body_start + 1 : body_end].strip()

    # 检查函数签名中是否包含 "unsafe"
    if "unsafe" in signature:
        # 去掉函数签名中的 "unsafe"
        signature = signature.replace("unsafe ", "")
        # 将 "unsafe" 移到函数体的最外层块
        body = f"unsafe {{\n{body}\n}}"

    # 重新组合函数
    new_rust_code = f"{signature} {{\n{body}\n}}"
    return new_rust_code

def first_round_optimize(path):
    PROJ_DIR = path
    SRC_DIR = os.path.join(PROJ_DIR, "src")
    TEST_DIR = os.path.join(PROJ_DIR, "tests")

    for root, dirs, files in os.walk(SRC_DIR):
        for file in files:
            logging.info(f"{file}第一轮优化")
            file_path = os.path.join(root, file)
            functions = split_rust_functions_in_file(file_path)
            if functions == None:
                logging.info(f"文件{file_path}第一轮优化失败，functions提取为空")
                return
            with open(file_path, 'r') as file:
                file_content = file.read()
            for func in functions:
                func_code = func['code']
                # func_code_lines = func_code.split('\n')

                optimized_func_code = transform_unsafe_fn(func_code)
                file_content = file_content.replace(func_code, optimized_func_code)
            with open(file_path, 'w') as file:
                file.write(file_content)

    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            logging.info(f"{file}第一轮优化")
            file_path = os.path.join(root, file)
            functions = split_rust_functions_in_file(file_path)
            if functions == None:
                logging.info(f"文件{file_path}第一轮优化失败，functions提取为空")
                return
            with open(file_path, 'r') as file:
                file_content = file.read()
            for func in functions:
                func_code = func['code']
                # func_code_lines = func_code.split('\n')

                optimized_func_code = transform_unsafe_fn(func_code)
                file_content = file_content.replace(func_code, optimized_func_code)
            with open(file_path, 'w') as file:
                file.write(file_content)

code = """
pub unsafe extern "C" fn test_slist_to_array() {
    let mut list: *mut SListEntry = 0 as *mut SListEntry;
    let mut array: *mut *mut libc::c_void = 0 as *mut *mut libc::c_void;
    list = generate_list();
    array = slist_to_array(list);
    assert!(
        *array.offset(0 as libc::c_int as isize)
            == &mut variable1 as *mut libc::c_int as *mut libc::c_void,
        "array[0] == &variable1"
    );
    assert!(
        *array.offset(1 as libc::c_int as isize)
            == &mut variable2 as *mut libc::c_int as *mut libc::c_void,
        "array[1] == &variable2"
    );
    assert!(
        *array.offset(2 as libc::c_int as isize)
            == &mut variable3 as *mut libc::c_int as *mut libc::c_void,
        "array[2] == &variable3"
    );
    assert!(
        *array.offset(3 as libc::c_int as isize)
            == &mut variable4 as *mut libc::c_int as *mut libc::c_void,
        "array[3] == &variable4"
    );
    alloc_test_free(array as *mut libc::c_void);
}
"""

if __name__ == "__main__":

    res = transform_unsafe_fn(code)
    print(res)